#!/usr/bin/env python3
"""
双阶段yolo识别+ocr兜底

ROS节点：集成GStreamer和YOLO检测，带结果缓存和投票机制

功能特点：
1. 订阅/yolo_detect话题，接收"yolostart"命令开始检测
2. 使用GStreamer接口获取视频帧
3. 使用双阶段YOLO管道进行目标检测
4. 为cone、number、dashboard分别维护结果缓存
5. 实现上进下出的缓存策略
6. 当所有缓存都满时，进行投票选择最终结果
7. 发布最终结果到/yolo_result话题
"""

import os
import sys
import time
from collections import Counter, deque
from typing import Deque, Optional, Tuple, Dict, Any
import threading

import rospy
from std_msgs.msg import String


def _resolve_paths() -> Tuple[str, str, str]:
    """解析项目路径"""
    # scripts -> dog_motion -> src -> nav_ws -> repo_root
    repo_root = os.path.dirname(
        os.path.dirname(
            os.path.dirname(
                os.path.dirname(
                    os.path.dirname(os.path.abspath(__file__))
                )
            )
        )
    )
    detect_dir = os.path.join(repo_root, 'yolo-jetson', 'detect')
    rs2_dir = os.path.join(repo_root, 'yolo-jetson', 'rs2')
    return repo_root, detect_dir, rs2_dir


REPO_ROOT, DETECT_DIR, RS2_DIR = _resolve_paths()
if DETECT_DIR not in sys.path:
    sys.path.append(DETECT_DIR)
if RS2_DIR not in sys.path:
    sys.path.append(RS2_DIR)

# 导入YOLO检测管道和GStreamer接口
try:
    from double_yolo_pipeline_with_cnocr import create_double_yolo_pipeline_with_cnocr
except ImportError:
    from yolo_pipe_ocr_hsv_final import create_ocr_hsv_pipeline as create_pipeline
    create_double_yolo_pipeline_with_cnocr = None

from GStreamerWrapper import GStreamerWrapper


class ResultCache:
    """结果缓存类，支持上进下出的队列操作"""
    
    def __init__(self, max_size: int = 5, name: str = "cache"):
        self.max_size = max_size
        self.name = name
        self.data: Deque[str] = deque(maxlen=max_size)
        self.lock = threading.Lock()
    
    def add(self, result: str) -> None:
        """添加结果到缓存顶部（最新位置）"""
        if result == "None":
            return
        
        with self.lock:
            # 如果缓存已满，会自动弹出最旧的元素
            self.data.append(result)
    
    def is_full(self) -> bool:
        """检查缓存是否已满"""
        with self.lock:
            return len(self.data) >= self.max_size
    
    def get_most_common(self) -> Optional[str]:
        """获取缓存中出现最多的结果"""
        with self.lock:
            if not self.data:
                return None
            counter = Counter(self.data)
            most_common_result, _ = counter.most_common(1)[0]
            return most_common_result
    
    def clear(self) -> None:
        """清空缓存"""
        with self.lock:
            self.data.clear()
    
    def size(self) -> int:
        """获取当前缓存大小"""
        with self.lock:
            return len(self.data)
    
    def to_list(self) -> list:
        """获取缓存内容的副本"""
        with self.lock:
            return list(self.data)


class YOLOGStreamerCacheNode:
    """YOLO GStreamer缓存节点"""
    
    def __init__(self) -> None:
        rospy.init_node('yolo_gstreamer_cache_node', anonymous=False)
        
        # 基础参数
        self.width = rospy.get_param('~width', 640)
        self.height = rospy.get_param('~height', 360)
        self.first_stage_conf = rospy.get_param('~first_stage_conf', 0.6)
        self.second_stage_conf = rospy.get_param('~second_stage_conf', 0.5)
        self.final_conf = rospy.get_param('~final_conf', 0.6)
        self.top_cut_ratio = rospy.get_param('~top_cut_ratio', 0.25)
        self.width_center_ratio = rospy.get_param('~width_center_ratio', 0.5)
        
        # 缓存参数
        self.cache_size = int(rospy.get_param('~cache_size', 5))
        self.target_fps = int(rospy.get_param('~target_fps', 10))
        self.debug_mode = rospy.get_param('~debug', True)
        
        # 模型路径发现
        self.model_path = self._find_model_path()
        if not self.model_path:
            rospy.logerr('未找到模型权重文件。请设置~model_path参数为有效的.pt文件路径。')
            raise RuntimeError('model_path not found')
        
        # 初始化YOLO检测管道
        self._init_pipeline()
        
        # 初始化GStreamer
        try:
            self.gs = GStreamerWrapper(width=self.width, height=self.height)
            rospy.loginfo(f'GStreamer初始化成功，分辨率: {self.width}x{self.height}')
        except Exception as e:
            rospy.logerr(f'GStreamer初始化失败: {e}')
            raise
        
        # 初始化结果缓存
        self.cone_cache = ResultCache(self.cache_size, "cone")
        self.number_cache = ResultCache(self.cache_size, "number") 
        self.dashboard_cache = ResultCache(self.cache_size, "dashboard")
        
        # ROS 发布/订阅
        self.result_pub = rospy.Publisher('/yolo_result', String, queue_size=10)
        self.cmd_sub = rospy.Subscriber('/yolo_detect', String, self._cmd_callback, queue_size=10)
        self.status_pub = rospy.Publisher('/yolo_status', String, queue_size=10)
        
        # 状态管理
        self._running = False
        self._frame_interval = 1.0 / max(1, self.target_fps)
        self._last_time = 0.0
        self._frame_count = 0
        self._success_count = 0
        self._processing_lock = threading.Lock()
        
        # 预热引擎
        self._warmup_engine()
        
        # 启动定时器
        self._timer = rospy.Timer(rospy.Duration(0.05), self._on_timer)  # 20Hz检查频率
        
        rospy.loginfo('YOLO GStreamer缓存节点已就绪')
        rospy.loginfo(f'模型: {self.model_path}')
        rospy.loginfo(f'分辨率: {self.width}x{self.height}')
        rospy.loginfo(f'目标FPS: {self.target_fps}')
        rospy.loginfo(f'缓存大小: {self.cache_size}')
        rospy.loginfo('发送"yolostart"到/yolo_detect话题开始检测')
    
    def _find_model_path(self) -> str:
        """查找模型权重文件"""
        model_path = rospy.get_param('~model_path', '')
        if model_path and os.path.exists(model_path):
            return model_path
        
        # 候选模型路径
        candidate_models = [
            os.path.join(DETECT_DIR, 'weights', 'all_yolo_all_epochs_0811', 'best.pt'),
            os.path.join(DETECT_DIR, 'weights', 'all_yolo_all_epochs_0808', 'best.pt'),
            os.path.join(DETECT_DIR, 'weights', 'all_yolo_all_epochs_0805', 'best.pt'),
            os.path.join(DETECT_DIR, 'weights', 'all_yolo_all_epochs_0804', 'best.pt'),
        ]
        
        for path in candidate_models:
            if os.path.exists(path):
                rospy.loginfo(f'找到模型: {path}')
                return path
        
        return ''
    
    def _init_pipeline(self) -> None:
        """初始化YOLO检测管道"""
        try:
            if create_double_yolo_pipeline_with_cnocr is not None:
                # 优先使用带cnOCR的管道
                self.pipeline = create_double_yolo_pipeline_with_cnocr(
                    model_path=self.model_path,
                    debug=self.debug_mode,
                    enable_cnocr=True
                )
                rospy.loginfo('使用双阶段YOLO + cnOCR管道')
            else:
                # 回退到标准管道
                self.pipeline = create_pipeline(
                    model_path=self.model_path,
                    debug=self.debug_mode
                )
                rospy.loginfo('使用标准YOLO OCR HSV管道')
        except Exception as e:
            rospy.logerr(f'YOLO管道初始化失败: {e}')
            raise
    
    def _warmup_engine(self) -> None:
        """预热YOLO引擎"""
        rospy.loginfo('正在预热YOLO引擎...')
        start_time = time.time()
        
        # 尝试获取真实帧进行预热
        frame = None
        while time.time() - start_time < 5.0:
            frame = self.gs.GetFrame()
            if frame is not None:
                break
            time.sleep(0.1)
        
        # 如果无法获取真实帧，使用虚拟帧
        if frame is None:
            try:
                import numpy as np
                frame = np.zeros((self.height, self.width, 3), dtype=np.uint8)
                rospy.logwarn('使用虚拟帧进行预热')
            except Exception as e:
                rospy.logerr(f'无法分配虚拟帧进行预热: {e}')
                raise
        
        try:
            _ = self.pipeline.process_image(
                frame,
                top_cut_ratio=self.top_cut_ratio,
                width_center_ratio=self.width_center_ratio,
                first_stage_conf=self.first_stage_conf,
                final_conf=self.final_conf,
            )
            rospy.loginfo('YOLO引擎预热完成')
        except Exception as e:
            rospy.logerr(f'YOLO引擎预热失败: {e}')
            raise
    
    def _cmd_callback(self, msg: String) -> None:
        """处理命令消息"""
        cmd = (msg.data or '').strip().lower()
        
        if cmd == 'yolostart' and not self._running:
            self._start_detection()
        elif cmd in ('stop', 'yolostop') and self._running:
            self._stop_detection()
        elif cmd == 'status':
            self._publish_status()
        elif cmd == 'clear':
            self._clear_caches()
    
    def _start_detection(self) -> None:
        """开始检测"""
        self._clear_caches()
        self._running = True
        self._frame_count = 0
        self._success_count = 0
        rospy.loginfo('开始YOLO检测')
        self._publish_status()
    
    def _stop_detection(self) -> None:
        """停止检测"""
        self._running = False
        rospy.loginfo('停止YOLO检测')
        self._publish_status()
    
    def _clear_caches(self) -> None:
        """清空所有缓存"""
        self.cone_cache.clear()
        self.number_cache.clear()
        self.dashboard_cache.clear()
        rospy.loginfo('清空所有结果缓存')
    
    def _publish_status(self) -> None:
        """发布状态信息"""
        status = {
            'running': self._running,
            'frame_count': self._frame_count,
            'success_count': self._success_count,
            'cone_cache_size': self.cone_cache.size(),
            'number_cache_size': self.number_cache.size(),
            'dashboard_cache_size': self.dashboard_cache.size(),
            'cone_cache_full': self.cone_cache.is_full(),
            'number_cache_full': self.number_cache.is_full(),
            'dashboard_cache_full': self.dashboard_cache.is_full(),
        }
        self.status_pub.publish(String(data=str(status)))
        rospy.loginfo(f'状态: {status}')
    
    def _on_timer(self, _event) -> None:
        """定时器回调函数，处理图像帧"""
        if not self._running or rospy.is_shutdown():
            return
        
        # 获取帧
        frame = self.gs.GetFrame()
        if frame is None:
            return
        
        # 控制帧率
        now = time.time()
        if now - self._last_time < self._frame_interval:
            return
        self._last_time = now
        
        # 处理帧（非阻塞）
        if self._processing_lock.acquire(blocking=False):
            try:
                self._process_frame(frame)
            finally:
                self._processing_lock.release()
    
    def _process_frame(self, frame) -> None:
        """处理单帧图像"""
        self._frame_count += 1
        
        try:
            # 运行YOLO检测
            result = self.pipeline.process_image(
                frame,
                top_cut_ratio=self.top_cut_ratio,
                width_center_ratio=self.width_center_ratio,
                first_stage_conf=self.first_stage_conf,
                second_stage_conf=self.second_stage_conf,
                final_conf=self.final_conf,
            )
            
            # 提取检测结果
            cone_color = result.get('cone_color', 'None')
            number = result.get('number', 'None')
            dashboard_state = result.get('dashboard_state', 'None')
            
            # 记录有效检测
            has_valid_result = False
            if cone_color != 'None':
                self.cone_cache.add(cone_color)
                has_valid_result = True
            if number != 'None':
                self.number_cache.add(number)
                has_valid_result = True
            if dashboard_state != 'None':
                self.dashboard_cache.add(dashboard_state)
                has_valid_result = True
            
            if has_valid_result:
                self._success_count += 1
                
            if self.debug_mode:
                rospy.loginfo(
                    f'帧 {self._frame_count}: '
                    f'cone={cone_color}, number={number}, dashboard={dashboard_state} | '
                    f'缓存大小: C={self.cone_cache.size()}, N={self.number_cache.size()}, D={self.dashboard_cache.size()}'
                )
            
            # 检查是否所有缓存都满了
            if (self.cone_cache.is_full() and 
                self.number_cache.is_full() and 
                self.dashboard_cache.is_full()):
                self._finalize_and_publish_results()
                
        except Exception as e:
            rospy.logerr(f'处理帧 {self._frame_count} 时出错: {e}')
    
    def _finalize_and_publish_results(self) -> None:
        """完成检测并发布最终结果"""
        # 获取每个缓存中出现最多的结果
        final_cone = self.cone_cache.get_most_common()
        final_number = self.number_cache.get_most_common() 
        final_dashboard = self.dashboard_cache.get_most_common()
        
        if final_cone and final_number and final_dashboard:
            # 构造最终结果字符串
            final_result = f"{final_cone}-{final_number}-{final_dashboard}"
            
            # 发布结果
            self.result_pub.publish(String(data=final_result))
            
            rospy.loginfo(f'发布最终结果: {final_result}')
            rospy.loginfo(f'缓存统计:')
            rospy.loginfo(f'  Cone: {self.cone_cache.to_list()}')
            rospy.loginfo(f'  Number: {self.number_cache.to_list()}')
            rospy.loginfo(f'  Dashboard: {self.dashboard_cache.to_list()}')
            rospy.loginfo(f'处理完成: {self._frame_count} 帧, {self._success_count} 成功检测')
        else:
            rospy.logwarn('无法生成完整的最终结果，某些缓存为空')
        
        # 重置状态
        self._clear_caches()
        self._running = False
        self._publish_status()
    
    def __del__(self):
        """析构函数"""
        try:
            if hasattr(self, 'gs'):
                self.gs.StopThread()
        except:
            pass


if __name__ == '__main__':
    try:
        node = YOLOGStreamerCacheNode()
        rospy.spin()
    except Exception as e:
        rospy.logerr(f'YOLO GStreamer缓存节点失败: {e}')
        import traceback
        traceback.print_exc()
